import openai

openai.api_key = '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'

# 装饰器函数
def openai_exception(func):
    def wrapper(*args, **kwargs):
        try:
            result = func(*args, **kwargs)
            return result
        except openai.error.APIError as e:
            #Handle API error here, e.g. retry or log
            print(f"OpenAI API returned an API Error: {e}")
            return f"OpenAI API returned an API Error: {e}"
        except openai.error.APIConnectionError as e:
            #Handle connection error here
            print(f"Failed to connect to OpenAI API: {e}")
            return f"Failed to connect to OpenAI API: {e}"
        except openai.error.RateLimitError as e:
            #Handle rate limit error (we recommend using exponential backoff)
            print(f"OpenAI API request exceeded rate limit: {e}")
            return f"OpenAI API request exceeded rate limit: {e}"
        except openai.error.AuthenticationError as e:
            print(f"OpenAI Incorrect API: {e}")
            return f"OpenAI Incorrect API: {e}"
    return wrapper

@openai_exception
def create_images(query, img_size):
    response = openai.Image.create(
        prompt= query,
        n=1,
        size=img_size,
        response_format='b64_json'
        )
    b64_json = response['data'][0]['b64_json']
    return 'data:image/png;base64,' + b64_json


@openai_exception
def fix_bug(query):
    """
    修改bug
    """
    query = f"##### Fix bugs in the below function\n \n### Buggy Python\n {query} \n    \n### Fixed Python"
    response = openai.Completion.create(
        model="text-davinci-003",
        prompt=query,
        temperature=0,
        max_tokens=1000,
        top_p=1.0,
        frequency_penalty=0.0,
        presence_penalty=0.0,
        stop=["###"]
    )
    return response.choices[0].text

@openai_exception
def chat(query):
    response = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": query},
            ]
    )
    return response.choices[0].message.content


# if __name__ == "__main__":
#     query = """
#     如何写一个爬虫
#     """
#     result = chat(query)
#     print(result)
